from neo4j.graph import Node

class QueryService:
    def __init__(self, neo4j_driver):
        self.driver = neo4j_driver

    def execute_query(self, query):
        """执行自定义查询"""
        with self.driver.session() as session:
            result = session.run(query)
            records = []
            for record in result:
                record_dict = {
                    key: dict(value) if isinstance(value, Node) else value 
                    for key, value in record.items()
                }
                records.append(record_dict)
            return records

    def get_all_relationship_data(self, limit=0):
        """获取所有关系数据"""
        queries = {
            'artifact-meaning': """
                MATCH (a:CulturalItem)-[r]-(m:Meaning) 
                RETURN a, r, m
            """,
            'artifact-story': """
                MATCH (a:CulturalItem)-[r]-(s:Story)
                WHERE s.story IS NOT NULL
                RETURN a, r, s
            """,
            'artifact-gene': """
                MATCH (a:CulturalItem)-[r]-(g:Gene)
                RETURN a, r, g
            """
            # 可以在这里添加更多的查询类型
        }

        all_results = {}
        with self.driver.session() as session:
            for query_type, query in queries.items():
                if limit > 0:
                    query += f" LIMIT {limit}"
                result = session.run(query)
                all_results[query_type] = self._process_relationship_result(result, query_type)

        return all_results

    def _process_relationship_result(self, result, query_type):
        """处理关系查询结果"""
        nodes, links = [], []
        seen_nodes = set()

        for record in result:
            source_node = record['a']
            
            # 添加源节点（对所有情况都需要）
            if source_node.element_id not in seen_nodes:
                nodes.append({
                    'element_id': source_node.element_id,
                    'labels': list(source_node.labels),
                    'properties': dict(source_node)
                })
                seen_nodes.add(source_node.element_id)

            # 处理普通关系
            target_key = 'm' if query_type == 'artifact-meaning' else ('s' if query_type == 'artifact-story' else 'g')
            if target_key in record and record[target_key] is not None:
                target_node = record[target_key]
                if target_node.element_id not in seen_nodes:
                    nodes.append({
                        'element_id': target_node.element_id,
                        'labels': list(target_node.labels),
                        'properties': dict(target_node)
                    })
                    seen_nodes.add(target_node.element_id)

                links.append({
                    'source': source_node.element_id,
                    'target': target_node.element_id,
                    'type': type(record['r']).__name__
                })

        return {'nodes': nodes, 'links': links}